import sys
import os
if os.path.exists("opt"):
    sys.path.append(".")
elif os.path.exists("../opt"):
    sys.path.append("..")

from opt.util.dataset import NeRFDataset
import argparse
import matplotlib.pyplot as plt

# main if
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=str, default="/mnt/data/wangziyi/data/nerf_synthetic/lego")
    parser.add_argument("--split", type=str, default="train")
    args = parser.parse_args()
    dataset = NeRFDataset(**vars(args))
    rays = dataset.rays
    dirs = rays.dirs
    origins = rays.origins

    points3D = dataset.c2w[:, :3, 3]
    distance = points3D.norm(dim=1)
    print(distance) # 2.6874 
    # plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points3D[:, 0], points3D[:, 1], points3D[:, 2])
    # show xyz
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()